// 每个线程块每步计算一个32x32的子矩阵，使用32x32的A和32x32的B

#include <assert.h>
#include <stdlib.h>

// CUDA runtime
#include "helper.h"
#include <cublas_v2.h>
#include <cuda_runtime.h>

#define ROW_MAJOR(row, col, ld) ((row) * (ld) + (col))
#define CEIL_DIV(x, y) (((x) + (y) - 1) / (y))

constexpr int BLOCK_SIZE = 32;    // TODO 修改BLOCK_SIZE到16， 32，试试看会发生什么？

// sBM = BN = BK，简化地址计算
constexpr int BM = BLOCK_SIZE;
constexpr int BN = BLOCK_SIZE;
constexpr int BK = BLOCK_SIZE;

// a = mxk, b = kxn
// a Shape: [(m/BM, BM), (k/BK, BK)]
// a Stride:[(BM*k, k ), (BK, 1)]
// a Index :  bx    tx    k1  ty

// b Shape: [(k/BK, BK), (n/BN, BN)]
// b Stride:[(BK*n, n), (BN, 1)]
// b Index :   k1  tx    by    ty

// c Shape: [(m/BM, BM), (n/BN, BN)]
// c Stride:[(BM*n, n),   (BN, 1)]
// c Index :  bx    tx    by    ty
__global__ void sgemm_tiling(int m, int n, int k, float *a, float *b, float *c) {
  // blockIdx control subpanel matrix

  const int tx = threadIdx.y;
  const int ty = threadIdx.x;
  const int bx = blockIdx.x;
  const int by = blockIdx.y;

  __shared__ float s_a[BM][BK];
  __shared__ float s_b[BK][BN];

  size_t load_a_smem_m = bx * BM + tx;
  size_t load_b_smem_n = by * BN + ty;
  float sum = 0.0f;

  for (size_t k1 = 0; k1 < CEIL_DIV(k, BK); k1 += 1) {
    size_t load_a_smem_k = k1 * BK + ty;
    size_t load_b_smem_k = k1 * BK + tx;
    size_t load_a_gmem_addr = ROW_MAJOR(load_a_smem_m, load_a_smem_k, k);
    size_t load_b_gmem_addr = ROW_MAJOR(load_b_smem_k, load_b_smem_n, n);
    s_a[tx][ty] = a[load_a_gmem_addr];
    s_b[tx][ty] = b[load_b_gmem_addr];
    __syncthreads();

    for (size_t k2 = 0; k2 < BK; k2++) {
      sum += s_a[tx][k2] * s_b[k2][ty];
    }
    __syncthreads();
  }
  size_t load_c_gmem_addr = ROW_MAJOR(bx * BM + tx, by * BN + ty, n);
  c[load_c_gmem_addr] = sum;
}

void SGEMM_GPU(cublasHandle_t handle, int m, int n, int k, float *d_A, float *d_B, float *d_C) {

  // dim3 block(BM, BN);
  dim3 block(BN, BM);
  dim3 grid(CEIL_DIV(m, BM), CEIL_DIV(n, BN));  

  sgemm_tiling<<<grid, block>>>(m, n, k, d_A, d_B, d_C);
}
